import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import gaussian_kde

# 示例数据：假设有3辆车，5个时间步，位置是随机生成的。
A = 3  # 车辆数量
T = 5  # 时间步长

# 假设车辆的位置数据如下 (格式为 [A, T, 2]，其中 2 是经纬度/坐标)
agent_positions = np.random.rand(A, T, 2)  # 生成随机位置数据，取值范围在[0, 1]

# 1. 将所有车辆在不同时间的位置信息提取出来（即提取所有的经纬度数据）
all_positions = agent_positions.reshape(-1, 2)  # 合并所有车辆的位置信息

# 2. 计算三维密度分布
# 我们将使用核密度估计（KDE）计算密度
kde = gaussian_kde([all_positions[:, 0], all_positions[:, 1]])

# 创建网格，用于计算密度
x = np.linspace(0, 1, 100)
y = np.linspace(0, 1, 100)
X, Y = np.meshgrid(x, y)
positions = np.vstack([X.ravel(), Y.ravel()])
Z = kde(positions).reshape(X.shape)

# 3. 绘制三维热力图
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制热力图，使用 'coolwarm' cmap 让高密度显示为红色，低密度显示为蓝色
surf = ax.plot_surface(X, Y, Z, cmap="coolwarm", edgecolor='none')

# 设置标题和标签
ax.set_title("3D Vehicle Spatial Probability Distribution")
ax.set_xlabel("X Coordinate")
ax.set_ylabel("Y Coordinate")
ax.set_zlabel("Density")

# 添加颜色条
fig.colorbar(surf, ax=ax, shrink=0.5, aspect=5)

# 显示图形
plt.show()
